import cv2 as cv
import sys
import os
from python_ai.common.xcommon import *
import matplotlib.pyplot as plt
import numpy as np
import time
import datetime


def my_show_img(img, title="no title", trans=None, **kwargs):
    global spn
    spn += 1
    plt.subplot(spr, spc, spn)
    if trans is not None:
        img = trans(img)
    plt.imshow(img, **kwargs)
    plt.axis('off')
    plt.title(title)


img_dir = '../../../../large_data/pic'
img_path1 = os.path.join(img_dir, 'opencv-logo-white.png')
img_path2 = os.path.join(img_dir, 'DSC05039.JPG')

spr = 2
spc = 6
spn = 0
plt.figure(figsize=[15, 6])

sep('load')
img_logo = cv.imread(img_path1, cv.IMREAD_UNCHANGED)
img_logo = cv.cvtColor(img_logo, cv.COLOR_BGRA2BGR)
img = cv.imread(img_path2, cv.IMREAD_UNCHANGED)
print(img_logo.shape)
LOGO_H, LOGO_W, _ = img_logo.shape
print(img.shape)

my_show_img(img_logo, 'img_logo', lambda x: cv.cvtColor(x, cv.COLOR_BGR2RGB))
my_show_img(img, 'img', lambda x: cv.cvtColor(x, cv.COLOR_BGR2RGB))

# ROI in tgt img
roi = img[0:LOGO_H, 0:LOGO_W]
my_show_img(roi, 'roi', lambda x: cv.cvtColor(x, cv.COLOR_BGR2RGB))

# mask
img_logo_gray = cv.cvtColor(img_logo, cv.COLOR_RGBA2GRAY)
my_show_img(img_logo_gray, 'img logo gray', cmap='gray')
ret, mask = cv.threshold(img_logo_gray, 10, 255, cv.THRESH_BINARY)
my_show_img(mask, 'mask', cmap='gray')

# mask inv
mask_inv = cv.bitwise_not(mask)
my_show_img(mask_inv, 'mask inv', cmap='gray')

# process logo
img_logo = cv.bitwise_and(img_logo, img_logo, mask=mask)
my_show_img(img_logo, 'img_logo masked by mask', lambda x: cv.cvtColor(x, cv.COLOR_BGR2RGB))

# process tgt
roi = cv.bitwise_and(roi, roi, mask=mask_inv)
my_show_img(roi, 'ori masked by mask_inv', lambda x: cv.cvtColor(x, cv.COLOR_BGR2RGB))

# add together
sep('add together')
print(img_logo.shape)
print(roi.shape)
dst = cv.add(img_logo, roi)
my_show_img(dst, 'dst', lambda x: cv.cvtColor(x, cv.COLOR_BGR2RGB))
img[0:LOGO_H, 0:LOGO_W] = dst
my_show_img(img, 'dst -> img', lambda x: cv.cvtColor(x, cv.COLOR_BGR2RGB))

plt.show()
